import torch

from my_zsl.tools.distributed_utils import MetricLogger, SmoothedValue, warmup_lr_scheduler


def _move_labels_to_device(labels, device):
    """将标签元组中的所有元素转移到目标设备"""
    return tuple(label.to(device) for label in labels)


def train_one_epoch(model, device, criterion, train_loader, optimizer, epoch, warmup=True):
    model.train()
    metric_logger = MetricLogger(delimiter="  ")
    metric_logger.add_meter('lr', SmoothedValue(window_size=1, fmt='{value:.6f}'))
    header = 'Epoch: [{}]'.format(epoch)

    lr_scheduler = None
    if epoch == 0 and warmup is True:  # 当训练第一轮（epoch=0）时，启用warmup训练方式，可理解为热身训练
        warmup_factor = 1.0 / 1000
        warmup_iters = min(1000, len(train_loader) - 1)

        lr_scheduler = warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor)

    mean_loss = torch.zeros(1).to(device)

    for step, [images, targets, _] in enumerate(metric_logger.log_every(train_loader, 50, header)):
        optimizer.zero_grad()
        images = images.to(device)
        targets = _move_labels_to_device(targets, device)
        load_labels, fault_labels, diameter_labels = targets

        outputs = model(images)

        loss, _ = criterion(images, outputs, load_labels, fault_labels, diameter_labels)

        mean_loss = (mean_loss * step + loss.detach()) / (step + 1)

        loss.backward()
        optimizer.step()

        if lr_scheduler is not None:  # 第一轮使用warmup训练方式
            lr_scheduler.step()

        metric_logger.update(loss=mean_loss)
        now_lr = optimizer.param_groups[0]["lr"]
        metric_logger.update(lr=now_lr)

    return mean_loss.item()


@torch.no_grad()
def valid_one_epoch(model, device, criterion, val_loader, epoch):
    model.eval()
    metric_logger = MetricLogger(delimiter="  ")
    header = 'Validation Epoch: [{}]'.format(epoch)
    val_loss = torch.zeros(1).to(device)
    result = {"y_pred": [], "y_true": [], 'val_loss': 0.}
    for step, [images, targets, labels] in enumerate(metric_logger.log_every(val_loader, 50, header)):
        images = images.to(device)
        targets = _move_labels_to_device(targets, device)
        load_labels, fault_labels, diameter_labels = targets
        outputs = model(images)
        loss, _ = criterion(images, outputs, load_labels, fault_labels, diameter_labels)
        val_loss = (val_loss * step + loss.detach()) / (step + 1)
        # 解析分类结果
        fault_pred = torch.argmax(torch.softmax(outputs['fault_class_pred'], dim=1), dim=1)
        cond_pred = torch.argmax(torch.softmax(outputs['condition_pred'], dim=1), dim=1)
        diam_pred = torch.argmax(torch.softmax(outputs['diameter_pred'], dim=1), dim=1)
        pred_tensor = torch.stack([fault_pred, cond_pred, diam_pred], dim=1)  # shape: [B, 3]

        # 2. 转成 Python list (每个元素是长度为 3 的 list)
        pred_list = pred_tensor.tolist()
        predicted = val_loader.dataset.find_key_by_label(pred_list)

        result["y_pred"].extend(predicted)
        result["y_true"].extend(labels.cpu().numpy())

    result['val_loss'] = val_loss.item()
    return result

    # for step, (images, labels) in enumerate(train_iterator):
    #     # 将数据转移到设备
    #     images, labels = images.to(device), labels.to(device)
    #     # 梯度清0
    #     optimizer.zero_grad()
    #     # 前向传播
    #     outputs = model(images)
    #     # 计算损失
    #     loss = criterion(outputs, labels)
    #     # 反向传播
    #     loss.backward()
    #     # 更新参数
    #     optimizer.step()
    #
    #     mean_loss = (mean_loss * step + loss.detach()) / (step + 1)
    #     _, predicted = torch.max(outputs, 1)
    #     # 设置进度条
    #     train_iterator.set_postfix(loss=loss.item(), mean_loss=mean_loss.item())
    #     all_predictions.extend(predicted.cpu().numpy())
    #     all_labels.extend(labels.cpu().numpy())
    #
    # train_accuracy = accuracy_score(y_true=all_labels, y_pred=all_predictions)
    # return mean_loss.item(), train_accuracy


# def find_key_by_label(reverse_dict, label_list):
#     """
#     通过标签列表查找对应的键
#     Args:
#         label_list: 标签列表，如[3, 3, 3]
#         reverse_dict
#     Returns:
#         对应的键（如"3-OR-021"），若不存在返回None
#     """
#     # 将列表转换为元组（可哈希，作为reverse_dict的键）
#     label_tuple = tuple(label_list)
#     return reverse_dict.get(label_tuple, None)  # 找不到时返回None

# if __name__ == '__main__':
#     from my_zsl.train_utils.dataset import create_dataloaders
#     from my_zsl.models.DecouplingNetwork import TripleDecouplingNetwork, TripleDecouplingLoss
#
#     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
#     train_loader, val_loader = create_dataloaders(r'D:\Code\2-ZSL\Zero-Shot-Learning\data\0HP\dataset',
#                                                   batch_size=128, train_shuffle=True)
#     model = TripleDecouplingNetwork().to(device)
#     criterion = TripleDecouplingLoss()
#     optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
#
#     train_one_epoch(model, device, criterion, train_loader, optimizer, 1)
#     res = valid_one_epoch(model, device, criterion, val_loader, 1)
#     train_accuracy = accuracy_score(y_true=res['y_true'], y_pred=res['y_pred'])
#     print(train_accuracy)
